import pandas as pd
import copy
import random
import itertools
from tqdm import tqdm
import numpy as np
from datetime import datetime

# 确保导入的是我们最新修改的版本
from strategy import Strategy
from backtest_engine import BacktestEngine
from data_handler import DataHandler

class Optimizer:
    """
    改进的参数优化器 - 适配新策略参数体系
    """
    def __init__(self, data_handler, param_grid):
        """
        初始化优化器。
        """
        self.data_handler = data_handler
        self.param_grid = param_grid
        self.best_params = None
        self.best_fitness = -float('inf')
        self.best_full_metrics = {}
        self.optimization_results = []
        
        # 多样性和收敛控制
        self.diversity_history = []
        self.no_improvement_generations = 0
        self.stagnation_threshold = 10  # 连续多少代没有改进就采取措施
        self.restart_threshold = 20     # 连续多少代没有改进就重启
        
        self.optimization_periods = self._generate_optimization_periods()
        
        print("--- 已生成以下优化周期 ---")
        for i, (start, end, period_type) in enumerate(self.optimization_periods):
            print(f"  周期 {i+1} ({period_type}): {start} -> {end}")
        print("--------------------------")

    def _generate_optimization_periods(self):
        """生成优化周期"""
        periods = []
        global_start_date = '2016-01-01'
        global_end_date = '2025-08-01'

        periods.append((global_start_date, global_end_date, 'full'))

        for year in range(2016, 2024): 
            start_date = f"{year}-01-01"
            target_end_date = f"{year + 2}-12-31" 
            end_date = min(target_end_date, global_end_date)
            periods.append((start_date, end_date, '3y'))

        for year in range(2016, 2020):
            start_date = f"{year}-01-01"
            target_end_date = f"{year + 6}-12-31"
            end_date = min(target_end_date, global_end_date)
            periods.append((start_date, end_date, '7y'))
        
        return periods

    def _calculate_fitness(self, params):
        """适应度计算"""
        cagr_full = []
        cagr_3y = []
        cagr_7y = []
        full_period_metrics = None

        for i, (start, end, period_type) in enumerate(self.optimization_periods):
            strategy = Strategy(params=params)
            engine = BacktestEngine(
                start_date=start,
                end_date=end,
                data_handler=self.data_handler,
                strategy=strategy
            )
            metrics, _, _ = engine.run_backtest()
            
            current_cagr = metrics['annualized_return_cagr'] if metrics and 'annualized_return_cagr' in metrics else -1.0
            
            if period_type == 'full':
                cagr_full.append(current_cagr)
            elif period_type == '3y':
                cagr_3y.append(current_cagr)
            elif period_type == '7y':
                cagr_7y.append(current_cagr)

            if period_type == 'full':
                full_period_metrics = metrics if metrics else {}

        full_cagr_val = np.mean(cagr_full) if cagr_full else -1.0
        avg_3y_cagr = np.mean(cagr_3y) if cagr_3y else -1.0
        avg_7y_cagr = np.mean(cagr_7y) if cagr_7y else -1.0

        final_fitness_components = [full_cagr_val, avg_3y_cagr, avg_7y_cagr]
        fitness = np.mean(final_fitness_components)
        
        return fitness, full_period_metrics

    def _calculate_population_diversity(self, population):
        """计算种群多样性"""
        if len(population) < 2:
            return 0.0
        
        total_distance = 0
        comparisons = 0
        
        for i in range(len(population)):
            for j in range(i + 1, len(population)):
                distance = self._calculate_individual_distance(population[i], population[j])
                total_distance += distance
                comparisons += 1
        
        return total_distance / comparisons if comparisons > 0 else 0.0

    def _calculate_individual_distance(self, ind1, ind2):
        """计算两个个体之间的距离"""
        distance = 0
        param_count = 0
        
        for key in ind1.keys():
            if key in ind2:
                if isinstance(ind1[key], (int, float)):
                    # 数值参数：标准化后计算欧几里得距离
                    if key in self.param_grid:
                        options = self.param_grid[key]
                        if len(options) > 1:
                            range_val = max(options) - min(options)
                            if range_val > 0:
                                normalized_dist = abs(ind1[key] - ind2[key]) / range_val
                                distance += normalized_dist ** 2
                                param_count += 1
                elif isinstance(ind1[key], list):
                    # 列表参数：计算元素差异比例
                    len1, len2 = len(ind1[key]), len(ind2[key])
                    max_len = max(len1, len2)
                    if max_len > 0:
                        diff_count = 0
                        for k in range(max_len):
                            val1 = ind1[key][k] if k < len1 else None
                            val2 = ind2[key][k] if k < len2 else None
                            if val1 != val2:
                                diff_count += 1
                        distance += (diff_count / max_len) ** 2
                        param_count += 1
                else:
                    # 其他类型：直接比较
                    distance += 1.0 if ind1[key] != ind2[key] else 0.0
                    param_count += 1
        
        return np.sqrt(distance / param_count) if param_count > 0 else 0.0

    def _initialize_population_advanced(self, population_size, seed_params=None):
        """高级种群初始化：使用分层采样确保多样性"""
        population = []
        
        # 1. 添加种子参数
        if seed_params:
            for seed in seed_params:
                fixed_seed = self._fix_threshold_constraint(copy.deepcopy(seed))
                population.append(fixed_seed)
        
        # 2. 生成角落解（极值组合）
        corner_solutions = self._generate_corner_solutions()
        for corner in corner_solutions:
            if len(population) < population_size:
                fixed_corner = self._fix_threshold_constraint(corner)
                population.append(fixed_corner)
        
        # 3. 分层采样填充剩余个体
        remaining = population_size - len(population)
        if remaining > 0:
            stratified_individuals = self._stratified_sampling(remaining)
            for ind in stratified_individuals:
                fixed_ind = self._fix_threshold_constraint(ind)
                population.append(fixed_ind)
        
        # 4. 确保种群大小并检查多样性
        while len(population) < population_size:
            params = self._generate_random_params()
            params = self._fix_threshold_constraint(params)
            
            # 检查是否与现有个体过于相似
            is_diverse = True
            for existing in population:
                if self._calculate_individual_distance(params, existing) < 0.1:
                    is_diverse = False
                    break
            
            if is_diverse or len(population) == 0:  # 防止无限循环
                population.append(params)
        
        print(f"初始种群多样性: {self._calculate_population_diversity(population):.4f}")
        return population[:population_size]

    def _generate_corner_solutions(self):
        """生成角落解（参数空间的极值组合）"""
        corner_solutions = []
        
        # 【更新】获取关键参数的极值（适配新参数体系）
        key_params = ['buy_threshold', 'sell_threshold', 'initial_buy_ratio', 'initial_sell_ratio',
                     'buy_base_factor', 'sell_base_factor']
        
        for param in key_params:
            if param in self.param_grid and len(self.param_grid[param]) >= 2:
                # 最小值组合
                min_params = self._generate_random_params()
                min_params[param] = min(self.param_grid[param])
                corner_solutions.append(min_params)
                
                # 最大值组合
                max_params = self._generate_random_params()
                max_params[param] = max(self.param_grid[param])
                corner_solutions.append(max_params)
        
        return corner_solutions

    def _stratified_sampling(self, n_samples):
        """分层采样生成多样化个体"""
        individuals = []
        
        # 将参数空间分层
        strata = max(1, int(np.sqrt(n_samples)))
        
        for i in range(n_samples):
            params = {}
            
            # 对每个参数进行分层采样
            for key, values in self.param_grid.items():
                if len(values) > 1:
                    # 将参数范围分成 strata 层，在每层中随机选择
                    layer_size = max(1, len(values) // strata)
                    layer_idx = (i % strata) * layer_size
                    layer_end = min(len(values), layer_idx + layer_size)
                    
                    if layer_idx < len(values):
                        layer_values = values[layer_idx:layer_end]
                        params[key] = random.choice(layer_values)
                    else:
                        params[key] = random.choice(values)
                else:
                    params[key] = values[0]
            
            # 验证参数
            params = self._validate_child(params)
            individuals.append(params)
        
        return individuals

    def _adaptive_tournament_selection(self, population_with_fitness, generation, max_generations):
        """自适应锦标赛选择：根据进化阶段调整选择压力"""
        # 早期：较小的锦标赛大小，保持多样性
        # 后期：较大的锦标赛大小，加强选择压力
        progress = generation / max_generations
        min_k, max_k = 2, 5
        k = int(min_k + (max_k - min_k) * progress)
        
        tournament_size = min(k, len(population_with_fitness))
        tournament_contenders = random.sample(population_with_fitness, tournament_size)
        
        # 在进化后期，偶尔使用适应度平方来增加选择压力
        if progress > 0.7 and random.random() < 0.3:
            return max(tournament_contenders, key=lambda x: x['fitness'] ** 2)['params']
        else:
            return max(tournament_contenders, key=lambda x: x['fitness'])['params']

    def _multi_crossover(self, parent1, parent2):
        """多种交叉策略随机选择"""
        strategy = random.choice(['uniform', 'parametric', 'adaptive'])
        
        if strategy == 'uniform':
            return self._uniform_crossover(parent1, parent2)
        elif strategy == 'parametric':
            return self._parametric_crossover(parent1, parent2)
        else:
            return self._adaptive_crossover(parent1, parent2)

    def _uniform_crossover(self, parent1, parent2):
        """均匀交叉"""
        child1, child2 = copy.deepcopy(parent1), copy.deepcopy(parent2)
        
        for key in parent1.keys():
            if random.random() < 0.5:
                child1[key], child2[key] = child2[key], child1[key]
        
        return self._validate_child(child1), self._validate_child(child2)

    def _parametric_crossover(self, parent1, parent2):
        """参数相关交叉：相关参数一起交换"""
        child1, child2 = copy.deepcopy(parent1), copy.deepcopy(parent2)
        
        # 【更新】定义参数组（适配新参数体系）
        param_groups = [
            ['buy_threshold', 'sell_threshold'],
            ['initial_buy_ratio', 'initial_sell_ratio'],
            ['buy_base_factor', 'sell_base_factor'],
            ['index_thresholds', 'max_ratios'],
            ['spread_factor', 'etf_factor']
        ]
        
        for group in param_groups:
            if random.random() < 0.5:
                for param in group:
                    if param in parent1 and param in parent2:
                        child1[param], child2[param] = child2[param], child1[param]
        
        return self._validate_child(child1), self._validate_child(child2)

    def _adaptive_crossover(self, parent1, parent2):
        """自适应交叉：基于参数类型选择交叉方式"""
        child1, child2 = copy.deepcopy(parent1), copy.deepcopy(parent2)
        
        for key in parent1.keys():
            if isinstance(parent1[key], list):
                # 列表参数：部分交换
                if random.random() < 0.5:
                    child1[key], child2[key] = self._crossover_lists(parent1[key], parent2[key])
            elif isinstance(parent1[key], (int, float)):
                # 数值参数：插值交叉
                if random.random() < 0.5:
                    alpha = random.random()
                    if key in self.param_grid:
                        # 在允许的参数值中找到最接近插值结果的值
                        interpolated = alpha * parent1[key] + (1 - alpha) * parent2[key]
                        options = self.param_grid[key]
                        closest = min(options, key=lambda x: abs(x - interpolated))
                        child1[key] = closest
                        child2[key] = parent2[key]
                    else:
                        child1[key], child2[key] = parent2[key], parent1[key]
            else:
                # 其他参数：简单交换
                if random.random() < 0.5:
                    child1[key], child2[key] = parent2[key], parent1[key]
        
        return self._validate_child(child1), self._validate_child(child2)

    def _crossover_lists(self, list1, list2):
        """列表参数的交叉操作"""
        if len(list1) == len(list2):
            # 长度相等：随机位置交换
            child1, child2 = list1[:], list2[:]
            crossover_point = random.randint(1, len(list1) - 1)
            child1[crossover_point:], child2[crossover_point:] = child2[crossover_point:], child1[crossover_point:]
            return child1, child2
        else:
            # 长度不等：随机选择一个
            return random.choice([list1, list2]), random.choice([list1, list2])

    def _adaptive_mutation(self, individual, base_mutation_rate, generation, max_generations):
        """自适应变异：根据进化阶段和多样性调整变异策略"""
        mutated_individual = copy.deepcopy(individual)
        
        # 根据进化阶段调整变异率
        progress = generation / max_generations
        current_diversity = self.diversity_history[-1] if self.diversity_history else 0.5
        
        # 多样性低时增加变异率
        diversity_factor = max(0.5, 2.0 - 4.0 * current_diversity)
        adjusted_mutation_rate = base_mutation_rate * diversity_factor
        
        # 进化后期降低变异率，但保持最小值
        stage_factor = max(0.3, 1.0 - 0.5 * progress)
        final_mutation_rate = adjusted_mutation_rate * stage_factor
        
        for key in mutated_individual.keys():
            if random.random() < final_mutation_rate and key in self.param_grid:
                mutated_individual = self._mutate_parameter(mutated_individual, key, progress)
        
        # 偶尔进行大变异（重新初始化部分参数）
        if random.random() < 0.05:  # 5%概率大变异
            n_params_to_reinit = random.randint(1, min(3, len(mutated_individual)))
            params_to_reinit = random.sample(list(mutated_individual.keys()), n_params_to_reinit)
            for param in params_to_reinit:
                if param in self.param_grid:
                    mutated_individual[param] = random.choice(self.param_grid[param])
        
        return self._validate_child(mutated_individual)

    def _mutate_parameter(self, individual, key, progress):
        """智能参数变异"""
        current_val = individual[key]
        options = self.param_grid[key]
        
        if isinstance(current_val, (int, float)) and len(options) > 2:
            # 数值参数：智能邻域搜索
            if random.random() < 0.7:  # 70%概率邻域变异
                if current_val in options:
                    current_idx = options.index(current_val)
                    # 根据进化阶段调整搜索范围
                    search_range = max(1, int(len(options) * 0.1 * (1 - progress)))
                    
                    # 选择邻域内的值
                    start_idx = max(0, current_idx - search_range)
                    end_idx = min(len(options), current_idx + search_range + 1)
                    neighbors = options[start_idx:end_idx]
                    individual[key] = random.choice(neighbors)
                else:
                    individual[key] = random.choice(options)
            else:
                # 30%概率完全随机
                individual[key] = random.choice(options)
        else:
            # 其他类型参数：标准随机选择
            individual[key] = random.choice(options)
        
        return individual

    def _inject_diversity(self, population_with_fitness, injection_ratio=0.2):
        """多样性注入：替换相似度高的个体"""
        n_inject = max(1, int(len(population_with_fitness) * injection_ratio))
        
        # 找到相似度最高的个体（排除精英）
        elite_size = 2
        candidates = population_with_fitness[elite_size:]
        
        # 计算每个个体与其他个体的平均距离
        avg_distances = []
        for i, ind1 in enumerate(candidates):
            total_distance = 0
            for j, ind2 in enumerate(candidates):
                if i != j:
                    total_distance += self._calculate_individual_distance(
                        ind1['params'], ind2['params']
                    )
            avg_distance = total_distance / (len(candidates) - 1) if len(candidates) > 1 else 0
            avg_distances.append((i + elite_size, avg_distance))
        
        # 选择距离最小的个体进行替换
        avg_distances.sort(key=lambda x: x[1])
        replace_indices = [idx for idx, _ in avg_distances[:n_inject]]
        
        # 生成新的多样化个体
        for idx in replace_indices:
            new_individual = self._generate_diverse_individual(population_with_fitness)
            new_individual = self._fix_threshold_constraint(new_individual)
            
            # 计算适应度并替换
            fitness, metrics = self._calculate_fitness(new_individual)
            population_with_fitness[idx] = {
                'params': new_individual,
                'fitness': fitness,
                'metrics': metrics
            }
        
        print(f"已注入 {n_inject} 个多样化个体")
        return population_with_fitness

    def _generate_diverse_individual(self, existing_population):
        """生成与现有种群差异较大的个体"""
        max_attempts = 50
        
        for attempt in range(max_attempts):
            new_individual = self._generate_random_params()
            
            # 检查与现有个体的最小距离
            min_distance = float('inf')
            for existing in existing_population:
                distance = self._calculate_individual_distance(
                    new_individual, existing['params']
                )
                min_distance = min(min_distance, distance)
            
            # 如果与现有个体足够不同，则接受
            if min_distance > 0.3:  # 距离阈值
                return new_individual
        
        # 如果无法找到足够不同的个体，返回随机个体
        return self._generate_random_params()

    def run_genetic_algorithm(self, n_generations=50, population_size=40, 
                              crossover_rate=0.7, mutation_rate=0.1, 
                              elitism_size=2, seed_params=None):
        """运行遗传算法：添加多样性维护和自适应机制"""
        print("\n--- 开始改进的遗传算法优化 ---")
        print(f"代数: {n_generations}, 种群大小: {population_size}")
        print(f"多样性维护: 开启, 自适应选择/变异: 开启")

        # 使用改进的初始化
        population = self._initialize_population_advanced(population_size, seed_params)
        progress_bar = tqdm(range(n_generations), desc="遗传算法进化进度")

        for generation in progress_bar:
            # 评估种群
            population_with_fitness = []
            for params in population:
                fitness, metrics = self._calculate_fitness(params)
                population_with_fitness.append({
                    'params': params,
                    'fitness': fitness,
                    'metrics': metrics
                })

            population_with_fitness.sort(key=lambda x: x['fitness'], reverse=True)

            # 更新最优解
            current_best_individual = population_with_fitness[0]
            if current_best_individual['fitness'] > self.best_fitness:
                self.best_fitness = current_best_individual['fitness']
                self.best_params = current_best_individual['params'].copy()
                self.best_full_metrics = current_best_individual['metrics']
                self.no_improvement_generations = 0
            else:
                self.no_improvement_generations += 1

            # 计算并记录种群多样性
            current_diversity = self._calculate_population_diversity([p['params'] for p in population_with_fitness])
            self.diversity_history.append(current_diversity)

            # 多样性维护措施
            if (self.no_improvement_generations >= self.stagnation_threshold or 
                current_diversity < 0.1):  # 多样性过低
                population_with_fitness = self._inject_diversity(population_with_fitness)
                self.no_improvement_generations = max(0, self.no_improvement_generations - 5)
                print(f"第 {generation+1} 代: 执行多样性注入")

            # 记录优化结果
            self.optimization_results.append({
                'generation': generation + 1,
                'best_fitness': self.best_fitness,
                'best_params': self.best_params,
                'best_full_metrics': self.best_full_metrics,
                'diversity': current_diversity
            })

            # 更新进度条
            cagr_display = self.best_full_metrics.get('annualized_return_cagr', 0)
            drawdown_display = self.best_full_metrics.get('max_drawdown', 0)
            progress_bar.set_postfix(
                best_avg_cagr=f"{self.best_fitness:.2%}",
                full_cagr=f"{cagr_display:.2%}",
                full_mdd=f"{drawdown_display:.2%}",
                diversity=f"{current_diversity:.3f}",
                stagnant=f"{self.no_improvement_generations}"
            )

            # 生成下一代
            next_generation = []
            
            # 精英保留
            for i in range(min(elitism_size, len(population_with_fitness))):
                next_generation.append(copy.deepcopy(population_with_fitness[i]['params']))

            # 生成剩余个体
            while len(next_generation) < population_size:
                # 自适应选择
                parent1 = self._adaptive_tournament_selection(population_with_fitness, generation, n_generations)
                parent2 = self._adaptive_tournament_selection(population_with_fitness, generation, n_generations)

                # 多种交叉策略
                if random.random() < crossover_rate:
                    child1, child2 = self._multi_crossover(parent1, parent2)
                else:
                    child1, child2 = copy.deepcopy(parent1), copy.deepcopy(parent2)
                
                # 自适应变异
                child1 = self._adaptive_mutation(child1, mutation_rate, generation, n_generations)
                child2 = self._adaptive_mutation(child2, mutation_rate, generation, n_generations)
                
                # 修复约束
                child1 = self._fix_threshold_constraint(child1)
                child2 = self._fix_threshold_constraint(child2)
                
                next_generation.append(child1)
                if len(next_generation) < population_size:
                    next_generation.append(child2)

            population = next_generation[:population_size]

        print("\n--- 改进的遗传算法优化完成 ---")
        print(f"最终种群多样性: {self.diversity_history[-1]:.4f}")
        print(f"最后改进代数: {n_generations - self.no_improvement_generations}")
        
        if self.best_params is None:
            print("警告: 未能找到任何有效的参数组合。")
        else:
            print("找到的最优参数全周期表现:")
            for key, value in self.best_full_metrics.items():
                if isinstance(value, float):
                    print(f"  {key}: {value:.4f}")
                else:
                    print(f"  {key}: {value}")

        return self.best_params, self.best_full_metrics

    def _fix_threshold_constraint(self, params):
        """修复阈值约束而非拒绝个体"""
        if params['sell_threshold'] >= params['buy_threshold']:
            mid_point = (params['sell_threshold'] + params['buy_threshold']) / 2
            gap = abs(params['buy_threshold'] - params['sell_threshold']) * 0.1
            
            if gap < 0.01:
                gap = 0.01
            
            params['sell_threshold'] = mid_point - gap
            params['buy_threshold'] = mid_point + gap
            
            if 'sell_threshold' in self.param_grid and 'buy_threshold' in self.param_grid:
                sell_options = self.param_grid['sell_threshold']
                buy_options = self.param_grid['buy_threshold']
                
                params['sell_threshold'] = min(sell_options, key=lambda x: abs(x - params['sell_threshold']))
                params['buy_threshold'] = min(buy_options, key=lambda x: abs(x - params['buy_threshold']))
                
                if params['sell_threshold'] >= params['buy_threshold']:
                    valid_combinations = [(s, b) for s in sell_options for b in buy_options if s < b]
                    if valid_combinations:
                        sell_val, buy_val = random.choice(valid_combinations)
                        params['sell_threshold'] = sell_val
                        params['buy_threshold'] = buy_val
        
        return params

    def _validate_child(self, child):
        """智能参数验证"""
        required_ratio_len = len(child['index_thresholds']) + 1
        if len(child['max_ratios']) != required_ratio_len:
            current_ratios = child['max_ratios']
            if len(current_ratios) < required_ratio_len:
                child['max_ratios'] = current_ratios + [current_ratios[-1]] * (required_ratio_len - len(current_ratios))
            else:
                child['max_ratios'] = current_ratios[:required_ratio_len]
            
            if 'max_ratios' in self.param_grid:
                valid_ratios = [r for r in self.param_grid['max_ratios'] if len(r) == required_ratio_len]
                if valid_ratios:
                    best_match = min(valid_ratios, 
                                   key=lambda x: sum(abs(a - b) for a, b in zip(child['max_ratios'], x)))
                    child['max_ratios'] = best_match
        
        return child

    def _generate_random_params(self):
        """【更新】生成随机参数（适配买卖独立因子）"""
        params = {}
        
        # 【更新】参数键列表，包含买卖独立的因子参数
        param_keys = [
            'buy_threshold', 'sell_threshold', 'initial_buy_ratio', 'initial_sell_ratio',
            'buy_base_factor', 'sell_base_factor', 
            'buy_spread_factor', 'sell_spread_factor',     # 买卖独立的股债利差因子
            'buy_etf_factor', 'sell_etf_factor',           # 买卖独立的ETF因子
            'max_trade_ratio', 'index_thresholds', 'max_ratios'
        ]
        
        for key in param_keys:
            if key in self.param_grid:
                params[key] = random.choice(self.param_grid[key])
        
        params = self._validate_child(params)
        return params